///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file hamiltonian.cpp
///
/// \brief   some implementation for the hamiltonian class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include<src/pcps.h>
#include<src/hamiltonian.h>
#include<src/mpi_interface.h>
#include<src/psi3_interface.h>
#include<src/molpro_interface.h>
#include<src/archive.h>
#include<src/timing.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   creates text files with the one and two electron integrals in dmrg input format
///
/// \param[in]   userinp  user input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template <class S> void pcps::Hamiltonian<S>::write_dmrg_ints(const pcps::Input & userinp) const {

  // get the file names for the one and two electron integrals
  const std::string int1_file = userinp.workdir() + "dmrg_int1.txt";
  const std::string int2_file = userinp.workdir() + "dmrg_int2.txt";

  // open the one electron integral file
  std::ofstream ofs_int1(int1_file.c_str());

  // check that the file is open
  if (!ofs_int1.is_open())
    throw pcps::Exception("failed to open file \"%s\" in pcps::Hamiltonian::write_dmrg_ints") % int1_file;

  // write the number of orbitals
  ofs_int1 << _norb << std::endl;

  // write the one electron integrals
  for (int p = 0; p < _norb; p++)
  for (int q = 0; q < _norb; q++)
    ofs_int1 << boost::format("%4i  %4i  %60s\n") % p % q % pcps::format_number("%22.12e", this->get_oei(p, q));

  // close the one electron integral file
  ofs_int1.close();

  // open the two electron integral file
  std::ofstream ofs_int2(int2_file.c_str());

  // check that the file is open
  if (!ofs_int2.is_open())
    throw pcps::Exception("failed to open file \"%s\" in pcps::Hamiltonian::write_dmrg_ints") % int2_file;

  // write the number of orbitals
  ofs_int2 << _norb << std::endl;

  // write the two electron integrals
  for (int p = 0; p < _norb; p++)
  for (int q = 0; q < _norb; q++)
  for (int r = 0; r < _norb; r++)
  for (int s = 0; s < _norb; s++)
    ofs_int2 << boost::format("%4i  %4i  %4i  %4i  %60s\n") % p % r % q % s % pcps::format_number("%22.12e", this->get_tei(p, q, r, s));

  // close the two electron integral file
  ofs_int2.close();

}

// explicitly instantiate the function
template void pcps::Hamiltonian< double               >::write_dmrg_ints(const pcps::Input &) const;
template void pcps::Hamiltonian< std::complex<double> >::write_dmrg_ints(const pcps::Input &) const;

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   prepares the Hubbard model Hamiltonian for fermions or spinless fermions
///
/// \param[in]   userinp  user input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template <class S> void pcps::Hamiltonian<S>::initialize_hubbard(const pcps::Input & userinp) {

  throw pcps::Exception("not yet implemented");

}

template void pcps::Hamiltonian< double               >::initialize_hubbard(const pcps::Input &);
template void pcps::Hamiltonian< std::complex<double> >::initialize_hubbard(const pcps::Input &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   computes the modified one electron integrals
///
///   The quantum chemistry electronic hamiltonian can be written as
///
///   \f[ H = \sum_{i j} h_{i j} a^\dag_i a_j + \frac{1}{2} \sum_{ p q r s} [pq|rs] a^\dag_p a^\dag_r a_s a_q \f]
///
///   where \f$h_{i j}\f$ are the one electron integrals stored in _oei and
///   \f$[pq|rs]\f$ are the two electron integrals stored in _tei.
///
///   The local energy function in quantum monte carlo is more easily computed if we
///   rearrange the Hamiltonian as 
///
///   \f[ H = \sum_{i j} f_{i j} a^\dag_i a_j + \frac{1}{2} \sum_{ p q r s} [pq|rs] a^\dag_p a_q a^\dag_r a_s \f]
///
///   where we have commuted \f$a_q\f$ to be next to \f$a^\dag_p\f$ and absorbed the
///   result of the delta function with \f$a^\dag_r\f$ into the one electron integrals,
///   creating the modified one electron integrals
///
///   \f[  f_{i j} = h_{i j} - \frac{1}{2} \sum_p  [ip|pj]  \f]
///
///   These modified one electron integrals are stored in _mei.
///
/// \param[in]   userinp  user input options
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template <class S> void pcps::Hamiltonian<S>::initialize_mei(const pcps::Input & userinp) {

  // modify the one electron integrals
  for (int i = 0; i < _norb; i++)
  for (int j = 0; j < _norb; j++) {
    _mei[i*_norb+j] = _oei[i*_norb+j];
    for (int p = 0; p < _norb; p++)
      _mei[i*_norb+j] -= 0.5 * this->get_tei(i,p,p,j);
  }

}

template void pcps::Hamiltonian< double               >::initialize_mei(const pcps::Input &);
template void pcps::Hamiltonian< std::complex<double> >::initialize_mei(const pcps::Input &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   loads the quantum chemistry integrals and computes their cholesky decomposition
///
/// \param[in]   userinp  user input options
/// \param[in]   source   a string to specify where the integrals should be read from
///
///////////////////////////////////////////////////////////////////////////////////////////////////

template <class S> void pcps::Hamiltonian<S>::initialize_integrals(const pcps::Input & userinp,
                                                                   const std::string & source) {

  // set the number of orbitals
  _norb = userinp.n_act_orb();

  // get integrals on root node
  if ( pcps::mpi::rank() == 0 ) {

    // allocate storage for the integrals
    _oei.assign(_norb*_norb, pcps::zero<S>());
    _mei.assign(_norb*_norb, pcps::zero<S>());
    _tei.assign( _norb*(_norb+1) * ( _norb*(_norb+1) + 2 ) / 8, pcps::zero<S>());

    // read integrals on root process
    if      (source == "molpro")
      pcps::get_molpro_ints(userinp, _e_cor_nuc, &_oei.at(0), &_tei.at(0));
    else if (source == "psi3")
      pcps::get_psi3_ints(userinp, _e_cor_nuc, &_oei.at(0), &_tei.at(0));
    else
      throw pcps::Exception("unknown integral source \"%s\" in pcps::Hamiltonian::initialize_integrals") % source;

    //// send integrals to all processes
    //pcps::mpi::bcast(_e_cor_nuc);
    //pcps::mpi::bcast(_oei);
    //pcps::mpi::bcast(_tei);

    // initialize the modified one electron integrals
    this->initialize_mei(userinp);

    // compute the cholesky decomposition of the two electron integrals
    pcps::start_timer("cholesky decomposition");
    pcps::cholesky_decomp(*this, *this);
    pcps::stop_timer("cholesky decomposition");
    std::cout << boost::format("cholesky threshold of %.12f resulted in %i cholesky vectors")
                 % this->chol_thresh()
                 % _n_cholesky
              << std::endl
              << std::endl;

    // delete the two electron integrals
    _tei.clear();

    // write the hamiltonian to disk
    this->to_archive("pcps_ham_dump.bin");

  }

  // read the hamiltonian from disk (this also broadcasts it to all processes)
  this->from_archive("pcps_ham_dump.bin");

//  // print nuclear repulsion energy
//  std::cout << boost::format("enuc = %20.12f") % _e_cor_nuc << std::endl;
//  std::cout << std::endl;
//
//  // print one electron integrals
//  std::cout << boost::format("one electron integrals:") << std::endl;
//  std::cout << std::endl;
//  for (int i = 0; i < _norb; i++)
//  for (int j = 0; j < _norb; j++)
//    std::cout << boost::format(" %4i %4i %20.12f") % i % j % _oei[i*_norb+j] << std::endl;
//  std::cout << std::endl;
//
//  // print two electron integrals
//  std::cout << boost::format("two electron integrals:") << std::endl;
//  std::cout << std::endl;
//  for (int i = 0; i < _norb; i++)
//  for (int j = 0; j < _norb; j++)
//  for (int k = 0; k < _norb; k++)
//  for (int l = 0; l < _norb; l++)
//    std::cout << boost::format(" %4i %4i %4i %4i %20.12f") % i % j % k % l % this->get_tei(i,j,k,l) << std::endl;
//  std::cout << std::endl;
//
//  // print two electron integral diagonal
//  std::cout << boost::format("diagonal two electron integrals:") << std::endl;
//  std::cout << std::endl;
//  for (int i = 0; i < _norb; i++)
//  for (int j = 0; j < _norb; j++)
//    std::cout << boost::format(" %4i %4i %20.12f") % i % j % this->get_tei(i,j,i,j) << std::endl;
//  std::cout << std::endl;
//
//  for (int i = 0; i < _cholesky_vecs.size(); i++) {
//    if (i % (_norb*_norb) == 0)
//      std::cout << std::endl;
//    std::cout << boost::format("cholesky %6i %20.12f") % i % _cholesky_vecs.at(i) << std::endl;
//  }
//
//  // print two electron integral error
//  double error = 0.0;
//  for (int i = 0; i < _norb; i++)
//  for (int j = 0; j < _norb; j++)
//  for (int k = 0; k < _norb; k++)
//  for (int l = 0; l < _norb; l++) {
//    S x = this->get_tei(i,j,k,l);
//    for (int p = 0; p < _n_cholesky; p++)
//      x -= _cholesky_vecs.at(_norb*_norb*p + _norb*i + j) * _cholesky_vecs.at(_norb*_norb*p + _norb*k + l);
//    if (std::abs(x) > this->chol_thresh())
//      std::cout << boost::format("error in integral %4i %4i %4i %4i %20.12f") % i % j % k % l % this->get_tei(i,j,k,l)
//                << std::endl;
//    error += pcps::square_norm(x);
//  }
//  std::cout << std::endl;
//  error = std::sqrt(error);
//  std::cout << boost::format("total error in two electron integrals is %22.12e") % error << std::endl;

}

template void pcps::Hamiltonian< double               >::initialize_integrals(const pcps::Input &, const std::string &);
template void pcps::Hamiltonian< std::complex<double> >::initialize_integrals(const pcps::Input &, const std::string &);


template<class T> void pcps::Hamiltonian<T>::to_archive(pcps::Archive & ar) const {

  ar << this->_norb;
  ar << this->_e_cor_nuc;
  ar << this->_oei;
  ar << this->_mei;
  //ar << this->_tei;
  ar << this->_n_cholesky;
  ar << this->_cholesky_vecs;
  //ar << this->_cholesky_vecs_t;

}

template<class T> void pcps::Hamiltonian<T>::to_archive(const std::string & filename) const {

  // only the root process writes to disk
  if (pcps::mpi::rank() == 0) {
    pcps::BinaryFileArchive ar(filename, "wt");
    ar << *this;
  }

}

template<class T> void pcps::Hamiltonian<T>::from_archive(pcps::Archive & ar) {

  pcps::mpi::read_and_bcast(ar, this->_norb);
  pcps::mpi::read_and_bcast(ar, this->_e_cor_nuc);
  pcps::mpi::read_and_bcast(ar, this->_oei);
  pcps::mpi::read_and_bcast(ar, this->_mei);
  //pcps::mpi::read_and_bcast(ar, this->_tei);
  pcps::mpi::read_and_bcast(ar, this->_n_cholesky);
  pcps::mpi::read_and_bcast(ar, this->_cholesky_vecs);
  //pcps::mpi::read_and_bcast(ar, this->_cholesky_vecs_t);

  _cholesky_vecs_t.resize(_cholesky_vecs.size());
  for (int j = 0; j < _n_cholesky; j++)
    pcps::xcopy(_norb*_norb, &_cholesky_vecs.at(_norb*_norb*j), 1, &_cholesky_vecs_t.at(j), _n_cholesky);

}

template<class T> void pcps::Hamiltonian<T>::from_archive(const std::string & filename) {

  // root process reads from disk
  if (pcps::mpi::rank() == 0) {

    pcps::BinaryFileArchive ar(filename, "r");
    ar >> *this;

  // other processes get their data from root and so we just pass them an empty archive
  } else {

    pcps::BinaryArchive ar;
    ar >> *this;

  }

}

template<class T, class S> pcps::BinaryArchiveTemplate<T> & operator<<(pcps::BinaryArchiveTemplate<T> & ar, const pcps::Hamiltonian<S> & ham) {

  // root process writes to disk
  if (pcps::mpi::rank() == 0)
    ham.to_archive(ar);

  // return the archive
  return ar;

}

template pcps::BinaryArchiveTemplate<std::stringstream> & operator<<(pcps::BinaryArchiveTemplate<std::stringstream> &,
                                                                     const pcps::Hamiltonian<double> &);
template pcps::BinaryArchiveTemplate<std::stringstream> & operator<<(pcps::BinaryArchiveTemplate<std::stringstream> &,
                                                                     const pcps::Hamiltonian<std::complex<double> > &);
template pcps::BinaryArchiveTemplate<std::fstream> & operator<<(pcps::BinaryArchiveTemplate<std::fstream> &,
                                                                const pcps::Hamiltonian<double> &);
template pcps::BinaryArchiveTemplate<std::fstream> & operator<<(pcps::BinaryArchiveTemplate<std::fstream> &,
                                                                const pcps::Hamiltonian<std::complex<double> > &);

template<class T, class S> pcps::BinaryArchiveTemplate<T> & operator>>(pcps::BinaryArchiveTemplate<T> & ar, pcps::Hamiltonian<S> & ham) {

  // root process reads from disk
  if (pcps::mpi::rank() == 0) {

    ham.from_archive(ar);

  // other processes get their data from root and so we just pass them an empty archive
  } else {

    pcps::BinaryArchive empty;
    ham.from_archive(empty);

  }

  // return the archive
  return ar;

}

template pcps::BinaryArchiveTemplate<std::stringstream> & operator>>(pcps::BinaryArchiveTemplate<std::stringstream> &,
                                                                     pcps::Hamiltonian<double> &);
template pcps::BinaryArchiveTemplate<std::stringstream> & operator>>(pcps::BinaryArchiveTemplate<std::stringstream> &,
                                                                     pcps::Hamiltonian<std::complex<double> > &);
template pcps::BinaryArchiveTemplate<std::fstream> & operator>>(pcps::BinaryArchiveTemplate<std::fstream> &,
                                                                pcps::Hamiltonian<double> &);
template pcps::BinaryArchiveTemplate<std::fstream> & operator>>(pcps::BinaryArchiveTemplate<std::fstream> &,
                                                                pcps::Hamiltonian<std::complex<double> > &);

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief  Get the [pq,pq] two electron integrals.
///
/// \param[out]     pqpq_ints   on exit, the [pq,pq] integrals
///
///////////////////////////////////////////////////////////////////////////////////////////////////
template<class T> void pcps::Hamiltonian<T>::get_pqpq_ints(std::vector<T> & pqpq_ints) const {

  pqpq_ints.resize(_norb*_norb);
  for (int p = 0; p < _norb; p++)
  for (int q = 0; q < _norb; q++) {
    pqpq_ints[_norb*p+q] = pcps::xdotc(_n_cholesky, &_cholesky_vecs_t[_n_cholesky * ( _norb * p + q )], 1,
                                                    &_cholesky_vecs_t[_n_cholesky * ( _norb * p + q )], 1);
    //pqpq_ints[_norb*p+q] = pcps::zero<T>();
    //for (int u = 0; u < _n_cholesky; u++) {
    //  const T x = _cholesky_vecs_t[ _n_cholesky * ( _norb * p + q ) + u ];
    //  pqpq_ints[_norb*p+q] += x * x;
    //}
  }

}

template void pcps::Hamiltonian<double>::get_pqpq_ints(std::vector<double> &) const;
template void pcps::Hamiltonian<std::complex<double> >::get_pqpq_ints(std::vector<std::complex<double> > &) const;
